package com.tencent.wxcloudrun.service.impl;

import cn.hutool.core.util.IdUtil;
import com.alibaba.fastjson.JSONObject;
import com.tencent.wxcloudrun.model.Message;
import com.tencent.wxcloudrun.service.RobotService;
import com.tencent.wxcloudrun.util.JacksonUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;

import java.io.*;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.Comparator;
import java.util.List;
import java.util.Optional;
import java.util.regex.Matcher;
import java.util.regex.Pattern;


@Slf4j
@Service
public class RobotServiceImpl implements RobotService {

    @Value("${forward_url}")
    private String FORWARD_URL;


    @Override
    public ResponseEntity<StreamingResponseBody> chatCompletions(String requestBody) {
        // String conversationId = "3ffbedad-fb7e-441e-a736-b1b003264d0e";
        // final String FORWARD_URL = "https://zhenliving-stg1.pingan.com.cn/api/zyn-ops/customer/agent/v1";

        log.info("### requestBody:{}", requestBody);

        JSONObject requestJsonObject = JSONObject.parseObject(requestBody);
        Object metadata = requestJsonObject.get("metadata");

        JSONObject metadataObject = JSONObject.parseObject(metadata.toString());
        String conversationId = metadataObject.getString("conversation_id");
        String urlAddress = String.format("%s/conversations/%s/messages?streaming=true", FORWARD_URL, conversationId);

        log.info("### urlAddress:{}", urlAddress);

        String messages = requestJsonObject.getString("messages");
        List<Message> messageList = JacksonUtils.ofList(messages, Message.class);
        String messageContent = messageList.stream()
                .filter(msg -> StringUtils.equalsIgnoreCase("user", msg.getRole()))
                .map(Message::getContent)
                .sorted(Comparator.reverseOrder())
                .findFirst().orElseThrow(RuntimeException::new);

        log.info("### messageContent:{}", messageContent);
        String requestJson = String.format("{\"message\":{\"content\":\"%s\"}}", messageContent);

        return ResponseEntity.ok()
                .contentType(MediaType.TEXT_EVENT_STREAM)
                .body(outputStream -> {
                    try {
                        URL url = new URL(urlAddress);
                        HttpURLConnection connection = (HttpURLConnection) url.openConnection();
                        connection.setRequestMethod("POST");
                        connection.setRequestProperty("Accept", "text/event-stream");
                        connection.setRequestProperty("Connection", "keep-alive");
                        connection.setRequestProperty("Content-Type", "application/json");
                        connection.setDoOutput(true);
                        try (OutputStream os = connection.getOutputStream()) {
                            byte[] bytes = requestJson.getBytes(StandardCharsets.UTF_8);
                            os.write(bytes, 0, bytes.length);
                        }

                        try (InputStream inputStream = connection.getInputStream();
                             BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8))
                        ) {
                            final String randomUUID = IdUtil.fastSimpleUUID();

                            // long startTimeMillis = System.currentTimeMillis();
                            // String startJsonStr = "{\"id\":\"%s\",\"object\":\"chat.completion.chunk\",\"created\":%s,\"model\":\"gpt-4o-mini\", \"system_fingerprint\": \"fp_44709d6fcb\", \"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"\"},\"logprobs\":null,\"finish_reason\":null}]}\n\n";
                            // String startMessage = String.format(startJsonStr, randomUUID, startTimeMillis);
                            // outputStream.write(startMessage.getBytes(StandardCharsets.UTF_8));

                            // 依次转换出content
                            contentPipeline(bufferedReader, outputStream, randomUUID);

                            long tailTimeMillis = System.currentTimeMillis();
                            String tailJsonStr = "{\"id\":\"%s\",\"object\":\"chat.completion.chunk\",\"created\":%s,\"model\":\"gpt-4o-mini\", \"system_fingerprint\": \"fp_44709d6fcb\", \"choices\":[{\"index\":0,\"delta\":{},\"logprobs\":null,\"finish_reason\":\"stop\"}]}\n\n";
                            String tailMessage = String.format(tailJsonStr, randomUUID, tailTimeMillis);
                            outputStream.write(tailMessage.getBytes(StandardCharsets.UTF_8));

                            outputStream.flush();
                        }
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                });
    }


    private void contentPipeline(BufferedReader bufferedReader, OutputStream outputStream, String randomUUID) throws IOException {
        String regex = "(?<=^data:).+";
        Pattern pattern = Pattern.compile(regex);

        boolean begin = false;
        String temp = null;

        String line = null;
        while ((line = bufferedReader.readLine()) != null) {
            Matcher matcher = pattern.matcher(line);
            if (matcher.find()) {
                temp = matcher.group().trim();
            }
            if (begin && temp != null) {
                JSONObject jsonObject = JSONObject.parseObject(temp);
                Object code = jsonObject.get("code");
                if (code != null && StringUtils.equals("success", code.toString())) {

                    Optional.ofNullable(jsonObject.get("data")).ifPresent(data -> {
                        JSONObject dataObject = JSONObject.parseObject(data.toString());
                        Optional.ofNullable(dataObject.get("content")).ifPresent(content -> {
                            String contentString = content.toString();
                            log.info("### contentString:{}", contentString);
                            long currentTimeMillis = System.currentTimeMillis();
                            String resultJsonStr = "{\"id\":\"%s\",\"object\":\"chat.completion.chunk\",\"created\":%s,\"model\":\"gpt-4o-mini\", \"system_fingerprint\": \"fp_44709d6fcb\", \"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\",\"content\":\"%s\"},\"logprobs\":null,\"finish_reason\":null}]}\n\n";
                            String formatMessage = String.format(resultJsonStr, randomUUID, currentTimeMillis, contentString);

                            byte[] bytes = formatMessage.getBytes(StandardCharsets.UTF_8);
                            try {
                                outputStream.write(bytes);
                            } catch (IOException ignored) {
                            }
                        });
                    });
                }
                begin = false;
            }

            if (StringUtils.equalsIgnoreCase("event: chunk_message", line)) {
                temp = null;
                begin = true;
            }
        }

    }
}
